package com.hp.service.impl;

import com.hp.req.ChatMemoryReq;
import com.hp.vo.ChatMemoryVO;
import com.hp.service.ChatService;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.StreamingChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.image.ImageModel;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.util.*;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * @author HP
 */
@Slf4j
@Service
public class ChatServiceImpl implements ChatService {


    @Resource
    private ChatClient chatClient;

    @Resource
    private StreamingChatModel streamingChatModel;

    @Resource
    private ImageModel imageModel;

    // 存储对话
    List<Message> messages = new ArrayList<>();


    @Override
    public String normalCompletion(String message) {

        Prompt prompt = new Prompt(new UserMessage(message));
        return chatClient.prompt(prompt).call().chatResponse().getResult().getOutput().getContent();
    }


    @Override
    public Map<String, String> streamCompletion(String message) {

        StringBuilder fullContent = new StringBuilder();

        streamingChatModel.stream(new Prompt(message))
                .flatMap(chatResponse -> Flux.fromIterable(chatResponse.getResults()))
                .map(content -> content.getOutput().getContent())
                .doOnNext(fullContent::append)
                .last()
                .map(lastContent -> Map.of(message, fullContent.toString()))
                .block();
        return Map.of(message, fullContent.toString());
    }


    @Override
    public void image(String input, HttpServletResponse response) {
        ImagePrompt imagePrompt = new ImagePrompt(input);
        ImageResponse imageResponse = imageModel.call(imagePrompt);
        String imageUrl = imageResponse.getResult().getOutput().getUrl();

        try {
            URL url = new URL(imageUrl);
            InputStream in = url.openStream();

            response.setHeader("Content-Type", MediaType.IMAGE_PNG_VALUE);
            response.getOutputStream().write(in.readAllBytes());
            response.getOutputStream().flush();
        } catch (IOException e) {
            response.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
        }
    }

    @Override
    public String chatMemoryByList(String systemMessage, String userMessage) {

        messages.add(new SystemMessage(systemMessage));
        messages.add(new UserMessage(userMessage));
        //第一轮对话
        Prompt prompt = new Prompt(messages);

        ChatResponse response = chatClient.prompt(prompt).call().chatResponse();
        String content = response.getResult().getOutput().getContent();

        messages.add(new AssistantMessage(content));

        return content;
    }

    @Override
    public ChatMemoryVO chatMemory(ChatMemoryReq chatMemoryReq) {

        //对话记忆的唯一标识
        String conversantId = Objects.isNull(chatMemoryReq.getChatId()) ? UUID.randomUUID().toString() : chatMemoryReq.getChatId();

        ChatResponse response = chatClient
                .prompt()
                .system(chatMemoryReq.getSystemMessage())
                .user(chatMemoryReq.getUserMessage())
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, conversantId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .call()
                .chatResponse();
        String content = response.getResult().getOutput().getContent();
        log.info("content: {}", content);

        return ChatMemoryVO.builder().chatId(conversantId).content(content).build();
    }

}